// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/functions/remote/client/remote.h>

#include <melon/io/async/EventBase.h>
#include <pollux/expression/expr.h>
#include <pollux/expression/vector_function.h>
#include <pollux/functions/remote/client/thrift_client.h>
#include <pollux/functions/remote/if/get_serde.h>
#include <pollux/functions/remote/if/gen-cpp2/RemoteFunctionServiceAsyncClient.h>
#include <pollux/type/hive/hive_type_serializer.h>
#include <pollux/vector/vector_stream.h>

namespace kumo::pollux::functions {
    namespace {
        std::string serializeType(const TypePtr &type) {
            // Use hive type serializer.
            return type::hive::HiveTypeSerializer::serialize(type);
        }

        class RemoteFunction : public exec::VectorFunction {
        public:
            RemoteFunction(
                const std::string &functionName,
                const std::vector<exec::VectorFunctionArg> &inputArgs,
                const RemoteVectorFunctionMetadata &metadata)
                : functionName_(functionName),
                  location_(metadata.location),
                  thriftClient_(getThriftClient(location_, &eventBase_)),
                  serdeFormat_(metadata.serdeFormat),
                  serde_(getSerde(serdeFormat_)) {
                std::vector<TypePtr> types;
                types.reserve(inputArgs.size());
                serializedInputTypes_.reserve(inputArgs.size());

                for (const auto &arg: inputArgs) {
                    types.emplace_back(arg.type);
                    serializedInputTypes_.emplace_back(serializeType(arg.type));
                }
                remoteInputType_ = ROW(std::move(types));
            }

            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr &outputType,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                try {
                    applyRemote(rows, args, outputType, context, result);
                } catch (const PolluxRuntimeError &) {
                    throw;
                } catch (const std::exception &) {
                    context.setErrors(rows, std::current_exception());
                }
            }

        private:
            void applyRemote(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr &outputType,
                exec::EvalCtx &context,
                VectorPtr &result) const {
                // Create type and row vector for serialization.
                auto remoteRowVector = std::make_shared<RowVector>(
                    context.pool(),
                    remoteInputType_,
                    BufferPtr{},
                    rows.end(),
                    std::move(args));

                // Send to remote server.
                remote::RemoteFunctionResponse remoteResponse;
                remote::RemoteFunctionRequest request;
                request.throwOnError_ref() = context.throwOnError();

                auto functionHandle = request.remoteFunctionHandle_ref();
                functionHandle->name_ref() = functionName_;
                functionHandle->returnType_ref() = serializeType(outputType);
                functionHandle->argumentTypes_ref() = serializedInputTypes_;

                auto requestInputs = request.inputs_ref();
                requestInputs->rowCount_ref() = remoteRowVector->size();
                requestInputs->pageFormat_ref() = serdeFormat_;

                // TODO: serialize only active rows.
                requestInputs->payload_ref() = rowVectorToIOBuf(
                    remoteRowVector, rows.end(), *context.pool(), serde_.get());

                try {
                    thriftClient_->sync_invokeFunction(remoteResponse, request);
                } catch (const std::exception &e) {
                    POLLUX_FAIL(
                        "Error while executing remote function '{}' at '{}': {}",
                        functionName_,
                        location_.describe(),
                        e.what());
                }

                auto outputRowVector = IOBufToRowVector(
                    remoteResponse.result().value().payload().value(),
                    ROW({outputType}),
                    *context.pool(),
                    serde_.get());
                result = outputRowVector->childAt(0);

                if (auto errorPayload = remoteResponse.result().value().errorPayload()) {
                    auto errorsRowVector = IOBufToRowVector(
                        *errorPayload, ROW({VARCHAR()}), *context.pool(), serde_.get());
                    auto errorsVector =
                            errorsRowVector->childAt(0)->as_flat_vector<StringView>();
                    POLLUX_CHECK(errorsVector, "Should be convertible to flat vector");

                    SelectivityVector selectedRows(errorsRowVector->size());
                    selectedRows.applyToSelected([&](vector_size_t i) {
                        if (errorsVector->is_null_at(i)) {
                            return;
                        }
                        try {
                            throw std::runtime_error(errorsVector->value_at(i));
                        } catch (const std::exception &) {
                            context.setError(i, std::current_exception());
                        }
                    });
                }
            }

            const std::string functionName_;
            melon::SocketAddress location_;

            melon::EventBase eventBase_;
            std::unique_ptr<RemoteFunctionClient> thriftClient_;
            remote::PageFormat serdeFormat_;
            std::unique_ptr<VectorSerde> serde_;

            // Structures we construct once to cache:
            RowTypePtr remoteInputType_;
            std::vector<std::string> serializedInputTypes_;
        };

        std::shared_ptr<exec::VectorFunction> createRemoteFunction(
            const std::string &name,
            const std::vector<exec::VectorFunctionArg> &inputArgs,
            const core::QueryConfig & /*config*/,
            const RemoteVectorFunctionMetadata &metadata) {
            return std::make_unique<RemoteFunction>(name, inputArgs, metadata);
        }
    } // namespace

    void registerRemoteFunction(
        const std::string &name,
        std::vector<exec::FunctionSignaturePtr> signatures,
        const RemoteVectorFunctionMetadata &metadata,
        bool overwrite) {
        exec::registerStatefulVectorFunction(
            name,
            signatures,
            std::bind(
                createRemoteFunction,
                std::placeholders::_1,
                std::placeholders::_2,
                std::placeholders::_3,
                metadata),
            metadata,
            overwrite);
    }
} // namespace kumo::pollux::functions
